/*******************************************************************************
 * Copyright 2013 Mindaugas Idzelis
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 ******************************************************************************/
package incontrol.core;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.collect.Lists.newArrayList;
import static java.util.Collections.singletonList;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.TimeUnit;

import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;

public final class PartitionAlg {

	private PartitionAlg() {

	}

	public static class WeightComparator<T> implements Comparator<T> {
		private Weigher<T> w;

		public WeightComparator(Weigher<T> w) {
			this.w = w;
		}

		@Override
		public int compare(T o1, T o2) {
			int w1 = w.getWeight(o1);
			int w2 = w.getWeight(o2);
			if (w1 == w2)
				return 0;
			else if (w1 < w2)
				return -1;
			else
				return 1;
		}

	}

	public static <T extends Weighted> List<List<T>> ckk(List<T> list) {
		return ckk(list, null, false, -1, null);
	}

	public static <T extends Weighted> List<List<T>> ckk(List<T> list,
			long timeout, TimeUnit unit) {
		checkNotNull(unit);
		return ckk(list, null, false, timeout, unit);
	}

	public static <T> List<List<T>> ckk(List<T> list, Weigher<T> w) {
		Preconditions.checkNotNull(w);
		return ckk(list, w, true, -1, null);
	}

	public static <T> List<List<T>> ckk(List<T> list, Weigher<T> w,
			long timeout, TimeUnit unit) {
		Preconditions.checkNotNull(w);
		Preconditions.checkNotNull(unit);
		return ckk(list, w, true, timeout, unit);
	}

	private static <T> List<List<T>> ckk(List<T> list, Weigher<T> w,
			boolean useWeigher, long timeout, TimeUnit unit) {
		Preconditions.checkNotNull(list);
		if (list.isEmpty())
			return Collections.emptyList();

		List<KKNode> nodes = new ArrayList<KKNode>();
		for (T o : list) {
			KKNode node = new KKNode();
			node.leaf = true;
			node.value = useWeigher ? w.getWeight(o) : ((Weighted) o)
					.getWeight();
			node.object = o;
			nodes.add(node);
		}
		Collections.sort(nodes, Collections.reverseOrder());

		CKKNode root = new CKKNode(null, nodes);
		CKKNode cursor = root;

		List<List<T>> bestSolution = null;
		int bestFitness = Integer.MAX_VALUE;

		long expires = 0;
		if (unit != null) {
			expires = System.currentTimeMillis()
					+ TimeUnit.MILLISECONDS.convert(timeout, unit);
		}
		while (true) {
			// System.out.println(System.identityHashCode(cursor));
			if (expires > 0) {
				if (System.currentTimeMillis() > expires) {
					if (bestSolution != null) {
						System.out.println("timeout reached");
						break;
					}
				}
			}

			List<KKNode> newNodes = newArrayList(cursor.value);

			int head = cursor.headSum();
			int tail = cursor.tailSum();
			if (head > tail) {
				// pruning rule - head in one set, all others in other set

				List<List<T>> sol = solve(cursor.value);
				int fit = getSumDiff(sol, w);
				if (fit < bestFitness) {
					System.out.println("Found better solution: " + fit);
					bestSolution = sol;
					// printSol(bestSolution.solution);
					bestFitness = fit;
				}

				cursor.done = true;
				cursor.diff = null;
				cursor.sum = null;
				CKKNode parent = cursor.parent;
				// cursor.parent = null;
				cursor = parent;
				if (cursor == null)
					break;

			}
			if (!cursor.done && cursor.diff == null) {
				KKNode first = newNodes.remove(0);
				KKNode second = newNodes.remove(0);

				KKNode diff = new KKNode();
				diff.value = first.value - second.value;
				diff.p1 = first;
				diff.p2 = second;

				if (newNodes.isEmpty()) {
					newNodes.add(diff);

					cursor.diff = new CKKNode(cursor, newNodes);
					List<List<T>> sol = solve(cursor.diff.value);
					int fit = getSumDiff(sol, w);
					if (fit < bestFitness) {
						System.out.println("Found better solution: " + fit);
						bestSolution = sol;
						// printSol(bestSolution.solution);
						bestFitness = fit;
					}
					cursor = cursor.parent;
					if (cursor == null)
						break;
				} else {
					int binarySearch = Collections.binarySearch(newNodes, diff,
							Collections.reverseOrder());
					if (binarySearch < 0) {
						binarySearch = -binarySearch - 1;
					}
					newNodes = newArrayList(newNodes);
					newNodes.add(binarySearch, diff);
					cursor.diff = new CKKNode(cursor, newNodes);
				}
				cursor = cursor.diff;
			} else if (!cursor.done && cursor.sum == null) {
				KKNode first = newNodes.remove(0);
				KKNode second = newNodes.remove(0);

				boolean done = newNodes.isEmpty();

				KKNode sum = new KKNode();
				sum.value = first.value + second.value;
				sum.p1 = first;
				sum.p2 = second;
				sum.sum = true;
				newNodes.add(0, sum);

				// put the solution in the CKKNode... make sure to do the
				// sum
				// solution when sumDone

				cursor.sum = new CKKNode(cursor, newNodes);
				List<List<T>> solve = solve(cursor.sum.value);

				int fitness = getSumDiff(solve, w);
				if (fitness < bestFitness) {
					bestFitness = fitness;
					bestSolution = solve;
					System.out.println("Found better solution: " + bestFitness);
				}
				if (done) {
					cursor.done = true;
					cursor.diff = null;
					cursor.sum = null;
					CKKNode parent = cursor.parent;
					// cursor.parent = null;
					cursor = parent;
					if (cursor == null)
						break;
				} else {
					cursor = cursor.sum;
				}
			} else {
				CKKNode parent = cursor.parent;
				cursor.parent = null;
				cursor.value = null;
				cursor = parent;
				if (cursor == null)
					break;
			}

			if (bestFitness == 0)
				break;

		}

		return bestSolution;

	}

	private static <T> int getSumDiff(List<List<T>> solve, Weigher<T> w) {
		int largest = Integer.MIN_VALUE;
		int smallest = Integer.MAX_VALUE;
		for (List<T> list : solve) {
			int sum = sum(list, w);
			if (sum < smallest) {
				smallest = sum;
			}
			if (sum > largest) {
				largest = sum;
			}
		}
		return largest - smallest;
	}

	private static <T> int sum(List<T> list, Weigher<T> w) {
		int totalSize = 0;
		for (int i = 0; i < list.size(); i++) {
			if (w == null) {
				totalSize += ((Weighted) list.get(i)).getWeight();
			} else {
				totalSize += w.getWeight(list.get(i));
			}
		}
		return totalSize;
	}

	public static <T extends Weighted> List<List<T>> greedy(List<T> list) {
		return greedy(list, null, false);
	}

	public static <T> List<List<T>> greedy(List<T> list, Weigher<T> w) {
		Preconditions.checkNotNull(w);
		return greedy(list, w, true);
	}

	private static <T> List<List<T>> greedy(List<T> list, Weigher<T> w,
			boolean useWeigher) {
		Preconditions.checkNotNull(list);
		List<T> copy = newArrayList(list);
		Collections.sort(copy, Collections.reverseOrder());
		List<List<T>> partitions = newArrayList();
		partitions.add(Lists.<T> newArrayList());
		partitions.add(Lists.<T> newArrayList());

		while (!copy.isEmpty()) {
			T a = copy.remove(0);
			if (sum(partitions.get(0), w) < sum(partitions.get(1), w)) {
				partitions.get(0).add(a);
			} else {
				partitions.get(1).add(a);
			}
		}
		return partitions;
	}

	public static <T> List<List<T>> greedy(List<T> list, int k) {
		return greedy(list, null, false);
	}

	public static <T> List<List<T>> greedy(List<T> list, int k, Weigher<T> w) {
		Preconditions.checkNotNull(w);
		return greedy(list, k, w, true);
	}

	private static <T> List<List<T>> greedy(List<T> list, int k, Weigher<T> w,
			boolean useWeigher) {
		Preconditions.checkNotNull(list);
		List<T> copy = newArrayList(list);
		Collections.sort(copy, Collections.reverseOrder());
		List<List<T>> partitions = newArrayList();
		for (int i = 0; i < k; i++) {
			partitions.add(Lists.<T> newArrayList());
		}

		while (!copy.isEmpty()) {
			T a = copy.remove(0);
			int smallest = Integer.MAX_VALUE;
			int smallIndex = Integer.MAX_VALUE;
			for (int i = 0; i < partitions.size(); i++) {
				List<T> p = partitions.get(i);
				int sum = sum(p, w);
				if (sum < smallest) {
					smallest = sum;
					smallIndex = i;
				}
			}
			partitions.get(smallIndex).add(a);
		}
		return partitions;
	}

	public static <T> List<List<Object>> kk(List<T> list) {
		Preconditions.checkNotNull(list);
		return kk(list, null, false);
	}

	public static <T> List<List<Object>> kk(List<T> list, Weigher<T> w) {
		Preconditions.checkNotNull(w);
		return kk(list, w, true);
	}

	private static <T> List<List<Object>> kk(List<T> list, Weigher<T> w,
			boolean useWeigher) {
		Preconditions.checkNotNull(list);
		List<KKNode> nodes = new ArrayList<KKNode>();
		for (T o : list) {
			KKNode node = new KKNode();
			node.leaf = true;
			node.value = useWeigher ? w.getWeight(o) : ((Weighted) o)
					.getWeight();
			node.object = o;
			nodes.add(node);
		}
		Collections.sort(nodes, Collections.reverseOrder());

		while (true) {
			KKNode first = nodes.remove(0);
			KKNode second = nodes.remove(0);

			KKNode node = new KKNode();
			node.value = first.value - second.value;
			node.p1 = first;
			node.p2 = second;
			if (nodes.isEmpty()) {
				nodes.add(node);
				break;
			} else {
				int binarySearch = Collections.binarySearch(nodes, node,
						Collections.reverseOrder());
				if (binarySearch < 0) {
					binarySearch = -binarySearch - 1;
				}
				nodes.add(binarySearch, node);
			}
		}
		// reconstruct tree using binary iteration
		KKNode root = nodes.get(0);
		List<List<Object>> sets = solve(singletonList(root));
		return sets;
	}

	public static <T> List<List<T>> kk(List<T> list, int k) {
		return kk(list, k, null, false);
	}

	public static <T> List<List<T>> kk(List<T> list, int k, Weigher<T> w) {
		Preconditions.checkNotNull(w);
		return kk(list, k, w, true);
	}

	private static <T> List<List<T>> kk(List<T> list, int k, Weigher<T> w,
			boolean useWeigher) {
		List<KWNode> nodes = new ArrayList<KWNode>();
		for (T o : list) {
			int i = useWeigher ? w.getWeight(o) : ((Weighted) o).getWeight();
			KWNode node = new KWNode(i);
			node.set = true;
			node.object = o;
			nodes.add(node);
		}
		Collections.sort(nodes, Collections.reverseOrder());

		List<Partition> parts = newArrayList();
		for (int i = 0; i < list.size(); i++) {
			List<KWNode> values = newArrayList();
			values.add(nodes.remove(0));
			for (int j = 0; j < k - 1; j++) {
				values.add(new KWNode(0));
			}
			parts.add(new Partition(values));
		}
		List<List<T>> solution = newArrayList();
		while (true) {

			Partition first = parts.remove(0);
			Partition second = parts.remove(0);

			if (parts.isEmpty()) {

				int i = 0;
				int j = k - 1;
				for (; i < k; i++, j--) {
					List<T> set = newArrayList();
					Iterables.addAll(set, Iterables.concat(first.value.get(i)
							.<T> getLeafsObjects(), second.value.get(j)
							.<T> getLeafsObjects()));
					solution.add(set);
				}
				break;
			} else {

				Partition p = merge(first, second, k);

				if (0 == p.getEmptySets()) {
					// its full, normalize
					p = p.normalize();
				}

				int binarySearch = Collections.binarySearch(parts, p,
						Collections.reverseOrder());
				if (binarySearch < 0) {
					binarySearch = -binarySearch - 1;
				}
				parts.add(binarySearch, p);
			}

		}
		// reconstruct tree using binary iteration
		return solution;

	}

	private static Partition merge(Partition first, Partition second, int k) {
		if (first.getEmptySets() == 0 || second.getEmptySets() == 0) {
			List<KWNode> combined = newArrayList();
			int i = 0;
			int j = k - 1;
			// not all of them dummy, the largest set one...
			for (; i < k; i++, j--) {
				KWNode kwNode = new KWNode(first.value.get(i),
						second.value.get(j), first.value.get(i).value
								+ second.value.get(j).value);
				kwNode.set = true;
				combined.add(kwNode);
			}
			return new Partition(combined);
		}
		List<KWNode> list = newArrayList();
		int cap = 0;
		for (KWNode n : first.value) {
			if (!n.set) {
				break;
			}
			list.add(n);
			cap++;
		}
		for (KWNode n : second.value) {
			if (!n.set)
				break;
			list.add(n);
			cap++;
		}
		for (int i = 0; i < k - cap; i++) {
			list.add(new KWNode(0));
		}
		return new Partition(list);
	}

	@SuppressWarnings("unchecked")
	private static <T> List<List<T>> solve(List<KKNode> root) {
		List<T>[] sets = new List[2];
		sets[0] = new ArrayList<T>();
		sets[1] = new ArrayList<T>();
		int color[] = new int[1];
		// if (root.get(0).sum) {
		// color[0]=1;
		// }
		boolean sum = root.get(0).sum;
		for (int i = 0; i < root.size(); i++) {
			// if (sum && i >= 1) {
			// // The rest of the list, should all go in the other group
			// color[0] = 1;
			// }
			partition(sets, root.get(i), null, color, sum, i);
		}
		return Arrays.asList(sets[0], sets[1]);
	}

	@SuppressWarnings("unchecked")
	private static <T> void partition(List<T>[] sets, KKNode node,
			KKNode parent, int[] color, boolean rootIsSum, int i) {

		if (node.leaf) {
			if (rootIsSum) {
				if (i == 0) {
					sets[1].add((T) node.object);

				} else {
					if (parent == null) {
						// this is an override
						sets[0].add((T) node.object);
					} else {
						// this is an override

						sets[color[0]].add((T) node.object);
						if (parent == null || !parent.sum) {
							if (color[0] == 0) {
								color[0] = 1;
							} else {
								color[0] = 0;
							}
						}
					}
				}
			} else {
				sets[color[0]].add((T) node.object);

				if (parent == null || !parent.sum) {
					if (color[0] == 0) {
						color[0] = 1;
					} else {
						color[0] = 0;
					}
				}
			}

		} else {
			// need to obey left/right splits for non-leafs
			partition(sets, node.p1, node, color, rootIsSum, i);
			partition(sets, node.p2, node, color, rootIsSum, i);
			// // need to flip color after sum is processed
			if (i > 0 && node.sum) {
				if (color[0] == 0) {
					color[0] = 1;
				} else {
					color[0] = 0;
				}
			}
		}
	}

	static class CKKNode {

		List<KKNode> value = new ArrayList<KKNode>();
		CKKNode diff;
		CKKNode sum;
		CKKNode parent;
		boolean done;
		static int atomic;

		CKKNode(CKKNode parent, List<KKNode> value) {
			this.parent = parent;
			this.value = Arrays.asList(Iterables.toArray(value, KKNode.class));
			// System.out.println("Node "+(atomic++));
		}

		public int tailSum() {
			int tail = 0;
			for (int i = 1; i < value.size(); i++) {
				tail += value.get(i).value;
			}
			return tail;
		}

		public int headSum() {
			int sum = value.get(0).value;
			return sum;
		}

		@Override
		public String toString() {
			return Arrays.toString((KKNode[]) value.toArray(new KKNode[value
					.size()]));
		}
	}

	static class KKNode implements Comparable<KKNode> {

		int value;
		Object object;
		KKNode p1;
		KKNode p2;
		boolean leaf;
		boolean sum;

		@Override
		public int compareTo(KKNode o) {
			int s1 = value;
			int s2 = o.value;
			if (s1 == s2)
				return 0;
			else if (s1 < s2)
				return -1;
			else
				return 1;
		}

		@Override
		public String toString() {
			return Integer.toString(value);
		}

	}

	static class KWNode implements Comparable<KWNode> {

		int value;
		Object object;
		KWNode p1;
		KWNode p2;
		boolean set;

		public KWNode(int value) {
			this.value = value;
		}

		public KWNode(KWNode parent, int value) {
			this.p1 = parent;
			this.value = value;
		}

		public KWNode(KWNode p1, KWNode p2, int value) {
			this.p1 = p1;
			this.p2 = p2;
			this.value = value;
		}

		@Override
		public String toString() {
			if (!set) {
				return "*";
			}
			return Integer.toString(value) + "("
					+ Joiner.on(',').join(getLeafs()) + ")";
		}

		@Override
		public int compareTo(KWNode o) {
			int s1 = sum();
			int s2 = o.sum();
			if (s1 == s2)
				return 0;
			else if (s1 < s2)
				return -1;
			else
				return 1;
		}

		public int sum() {
			int sum = 0;
			for (int i : getLeafs()) {
				sum += i;
			}
			return sum;
		}

		public List<Integer> getLeafs() {
			List<Integer> leafs = newArrayList();
			List<KWNode> queue = newArrayList();
			queue.add(this);
			while (!queue.isEmpty()) {
				KWNode n = queue.remove(0);
				if (n.p1 == null && n.p2 == null) {
					if (n.set) {
						leafs.add(n.value);
					}
				} else {
					if (n.p1 != null)
						queue.add(n.p1);
					if (n.p2 != null)
						queue.add(n.p2);
				}

			}
			return leafs;
		}

		@SuppressWarnings("unchecked")
		public <T> List<T> getLeafsObjects() {
			List<T> leafs = newArrayList();
			List<KWNode> queue = newArrayList();
			queue.add(this);
			while (!queue.isEmpty()) {
				KWNode n = queue.remove(0);
				if (n.p1 == null && n.p2 == null) {
					if (n.set) {
						leafs.add((T) n.object);
					}
				} else {
					if (n.p1 != null)
						queue.add(n.p1);
					if (n.p2 != null)
						queue.add(n.p2);
				}

			}
			return leafs;
		}
	}

	static class Partition implements Comparable<Partition> {

		List<KWNode> value = new ArrayList<KWNode>();

		public Partition(List<KWNode> value) {
			this.value = value;
			Collections.sort(this.value, Collections.reverseOrder());
		}

		@Override
		public String toString() {
			return Arrays.toString((KWNode[]) value.toArray(new KWNode[value
					.size()]));
		}

		public int getEmptySets() {
			int empty = 0;
			for (KWNode n : value) {
				if (!n.set)
					empty++;
			}
			return empty;
		}

		public int min() {
			int min = Integer.MAX_VALUE;
			for (KWNode n : value) {
				if (n.value < min) {
					min = n.value;
				}
			}
			return min;
		}

		public Partition normalize() {
			int smallest = min();
			List<KWNode> norm = newArrayList();
			for (KWNode n : value) {
				KWNode node = new KWNode(n, n.value - smallest);
				node.set = true;
				norm.add(node);
			}
			return new Partition(norm);
		}

		public int sum() {
			int sum = 0;
			for (KWNode n : value) {
				sum += n.value;
			}
			return sum;
		}

		@Override
		public int compareTo(Partition o) {
			int s1 = sum();
			int s2 = o.sum();
			if (s1 == s2)
				return 0;
			else if (s1 < s2)
				return -1;
			else
				return 1;
		}
	}
}
